﻿#include "../../inc/helper.hpp"

const static int g_heart_beat_liveness = 3;
const static int g_heart_beat_interval = 1000;

const static char* g_msg_worker_ready = "\001";
const static char* g_msg_heart_beat = "\002";

struct worker_t
{
	zmq::message_t identity;
	std::string id_string;
	Time expiry;
};
using worker_list = std::list<worker_t>;
using zmsg_list = std::deque<zmq::message_t>;

worker_t worker_new(zmq::message_t& msg)
{
	worker_t wk;
	wk.identity.copy(&msg);
	wk.id_string = to_string(msg);
	wk.expiry = Now() + absl::Milliseconds(g_heart_beat_liveness * g_heart_beat_interval);
	return wk;
}

void worker_ready(worker_list& workers, worker_t&& worker)
{
	workers.remove_if([&worker](worker_t& wk) {
		return wk.id_string == worker.id_string;
	});
	workers.push_back(std::move(worker));
}

void worker_purge(worker_list& workers)
{
	workers.remove_if([](worker_t& wk) {
		return wk.expiry < Now();
	});
}

int main()
{
	zmq::context_t ctx;
	zmq::socket_t frontend(ctx, zmq::socket_type::router);
	zmq::socket_t backend(ctx, zmq::socket_type::router);
	frontend.bind("tcp://*:5555");
	backend.bind("tcp://*:5556");

	Time heartbeat_at = Now() + absl::Milliseconds(g_heart_beat_interval);
	std::list<worker_t> worker_queue;
	while (true)
	{
		std::vector<zmq::pollitem_t> items = {
			{backend, 0, ZMQ_POLLIN, 0},
			{frontend, 0, ZMQ_POLLIN, 0}
		};

		if (worker_queue.empty()) {
			zmq::poll(items.data(), 1, g_heart_beat_interval);
		}
		else {
			zmq::poll(items, g_heart_beat_interval);
		}

		if (items[0].revents & ZMQ_POLLIN)
		{
			/// 获取worker id
			auto replies = zmsg_recv(backend);

			/// 添加到可用worker列表中
			worker_ready(worker_queue, worker_new(replies[0]));

			if (replies.size() == 2)
			{
				std::string msg = to_string(replies.back());
				if (msg != g_msg_heart_beat &&
					msg != g_msg_worker_ready)
				{
					info("invalid message from worker");
				}
			}
			else
			{
				/// 如果是消息回复，转发到前端
				replies.erase(replies.begin(), replies.begin() + 1);
				zmsg_send(frontend, replies);
			}
		}

		if (items[1].revents & ZMQ_POLLIN)
		{
			auto requests = zmsg_recv(frontend);

			std::deque<zmq::message_t> msgs(1);
			msgs[0].copy( &(worker_queue.front().identity) );
			for (auto& msg : requests)
			{
				msgs.emplace_back(std::move(msg));
			}

			zmsg_send(backend, msgs.begin(), msgs.end());
			
			worker_queue.pop_front();
		}

		if (Now() >= heartbeat_at)
		{
			heartbeat_at = Now() + absl::Milliseconds(g_heart_beat_interval);
			for (worker_t& wk : worker_queue)
			{
				zmsg_list heartbeats(2);
				heartbeats[0].copy(&wk.identity);
				heartbeats[1] = to_message(g_msg_heart_beat);
				zmsg_send(backend, heartbeats);
			}
		}

		worker_purge(worker_queue);
	}
}